#! /usr/bin/env python
#
# Copyright 2010 by Akkana Peck akkana@shallowsky.com
# ... share and enjoy under the GPLv2 or (at your option) later.
#

import sys, os
import subprocess
import re
import shutil

class AccessPoint :
    """ One Cell or AccessPoint from iwlist output"""

    def __init__(self) :
        self.clear()

    def clear(self) :
        self.address = ""
        self.essid = ""
        self.protocol = ""
        self.encryptionType = ""

    def encryptionString(self) :
        if self.encryptionType == 'WPA' :
            return 'WPA'
        elif self.encrypted :
            return "encrypted"
        else :
            return "open"

def print_current_scheme() :
    try :
        fp = open('/etc/network/schemes/current', 'r')
        line = fp.readline()
        print "Current scheme:", line,
    except IOError:
        print "Current scheme unknown"

def parse_iwlist_output(maxnum=10, interface="eth1") :
    proc = subprocess.Popen('iwlist scan 2>/dev/null', shell=True, stdout=subprocess.PIPE, )
    #proc = subprocess.Popen('cat /home/akkana/iwlist.out 2>/dev/null', shell=True, stdout=subprocess.PIPE, )
    stdout_str = proc.communicate()[0]
    stdout_list = stdout_str.split('\n')

    ap = None
    aplist=[]

    for line in stdout_list:
        line=line.strip()

        match = re.search('Cell ', line)
        if match :
            ap = AccessPoint()
            aplist.append(ap)

        match = re.search('ESSID:"(.+)"', line)
        if match :
            if match.group(1) == "<hidden>":
                ap.essid = None
            # I have no idea what these \x00\x00\x00\x00\x00 essids are,
            # but they're quite common, and annoying to see in a UI:
            elif match.group(1) == "\\x00\\x00\\x00\\x00\\x00":
                ap.essid = "[null]"
            else :
                ap.essid = match.group(1)

        match = re.search('Address: (\S+)', line)
        if match:
            ap.address = match.group(1)

        match = re.search('Encryption key:([onf]+)', line)
        if match:
            if match.group(1) == "off" :
                ap.encrypted = False
            else :
                ap.encrypted = True

        match = re.search('Protocol:IEEE(.+)', line)
        if match:
            ap.protocol = match.group(1)

        match = re.search('WPA', line)
        if match:
            ap.encryptionType = "WPA"

        match = re.search('Mode:(.+)', line)
        if match :
            ap.mode = match.group(1)

    return aplist

def list_accesspoints(maxap) :
    aplist = parse_iwlist_output(maxap, "eth1")

    print "Visible access points:"
    for ap in aplist :
        # Show mode if it's anything but Master -- e.g. Ad-hoc
        mode = ''
        if ap.mode != 'Master' :
            mode = ' (' + ap.mode + ')'
        print "%-19s : %-10s %-10s %-10s" % (ap.essid, ap.protocol,
                                       ap.encryptionString(), mode)
        continue

def get_existing_schemes() :
    schemes = []
    dirname = "/etc/network/schemes"
    files = os.listdir(dirname)
    for f in files :
        if "interfaces-" != f[0:11] :
            continue
        schemename = f[11:]
        fp = open(os.path.join(dirname, f), "r")
        while 1:
            line = fp.readline()
            if not line : break
            if 'wireless-essid ' == line[0:15] :
                essid = line[15:]
                essid = essid.strip()
                schemes.append( [schemename, essid] )
                continue
    return schemes

def list_schemes() :
    schemes = get_existing_schemes()
    width = 0
    for s in schemes :
        if len(s[0]) > width : width = len(s[0])
    width += 1
    for s in schemes :
        print '{0:{1}}: {2}'.format(s[0], width, s[1])

def reset_current_scheme() :
    ifupdown(False)
    ifupdown(True)
    os.system("service networking restart")

def ifupdown(up=True) :
    # Find out what interfaces are currently in the interfaces file:
    interfaces = []
    fp = open('/etc/network/interfaces', 'r')
    while 1:
        line = fp.readline()
        if not line : break
        line.strip()
        if line[0:6] != 'iface ' : continue
        line = line[6:]
        line.strip()
        space = line.find(' ')
        if space > 0 :
            line = line[0:space]
        interfaces.append(line)
    fp.close()
    print "Interfaces:", interfaces
    if up :
        updown = "up"
    else :
        updown = "down"
    for iface in interfaces :
        print "Running if" + updown + " " + iface
        os.system("if" + updown + " " + iface)
        print "Running ifconfig " + iface + " " + updown
        os.system("ifconfig " + iface + " " + updown)

    # It would be nice not to have to do this, or at least not to do
    # it using os.system. Haven't yet found a comparably reliable way.
    if up :
        os.system("service networking restart")

#
# Copy an existing scheme info into /etc/network/interfaces
# and /etc/resolv.conf, then enable/start them.
#
def set_scheme(newscheme) :
    # Check whether the named scheme exists:
    filename = "/etc/network/schemes/interfaces-" + newscheme

    # Take current interfaces down first:
    ifupdown(False)

    if not os.access(filename, os.R_OK) :
        ask_to_make_new_scheme(newscheme)
        return
    else :
        shutil.copy2(filename, "/etc/network/interfaces")

        # Copy a resolv.conf too, if any
        filename = "/etc/network/schemes/resolv.conf-" + newscheme
        if os.access(filename, os.R_OK) :
            print "Also copying", filename
            os.rename("/etc/resolv.conf", "/etc/resolv.conf.bak")
            shutil.copy2(filename, "/etc/resolv.conf")
        else: print "No", filename, "to copy"

    # Okay, now the scheme is in place. We need to tell the system to use it:
    ifupdown(True)

    # Update the current scheme file:
    fp = open('/etc/network/schemes/current', 'w')
    fp.write(newscheme + '\n')
    fp.close

#
# Prompt the user whether to create a new scheme, then do so if confirmed.
#
def ask_to_make_new_scheme(schemename) :
    print "No scheme named", schemename
    ans = raw_input("Use it anyway? (Y/n) ")
    if ans != 'y' and ans != 'Y' and ans != '':
        sys.exit(1)

    # Save, if desired:
    ans = raw_input("Save this scheme for later? (y/N) ")
    print "Read answer", ans
    if ans == 'y' or ans == 'Y' :
        save = True
    else :
        save = False

    create_new_interfaces_file(schemename, save)

#
# Create a new interfaces file with the new scheme,
# and also enables the interface
#
def create_new_interfaces_file(schemename, save) :
    # Take current interfaces down first:
    ifupdown(False)

    interface = 'eth1'        # XXX make configurable or figure it out
    interfacefile = '/etc/network/interfaces'
    fp = open(interfacefile, "w")
    fp.write('''auto lo
iface lo inet loopback

auto %s
allow-hotplug %s
iface %s inet dhcp
wireless-essid %s
wireless-key off
''' % (interface, interface, interface, schemename))
    #fp.write("wireless-key: " + $pw")
    fp.close()
    # Why set the key to off specifically? Because at least on Intel chips,
    # if the last connection used a restricted key, somehow the hardware
    # (or something else outside of the interfaces file) seems to remember
    # that even through reboots and power cycles, and won't work again
    # until the key is reset to off.

    # If a key *is* needed, you may need an additional line, such as:
    # pre-up iwconfig eth1 essid [essid] key restricted [passphrase]
    # unfortunately there doesn't seem to be a way to set a restricted
    # passphrase without calling iwconfig via pre-up -- you might think
    # wireless-key restricted [passphrase] might work, but it doesn't,
    # and using iwconfig to set the passphrase without specifying the
    # essid in the same line doesn't work either. Oy!
    # Even more unfortunate, it's not clear how to tell from iwlist output
    # which accesspoints do or don't need a restricted key.

    if save :
        shutil.copy2(interfacefile,
                     '/etc/network/schemes/interfaces-' + schemename)
    else :
        print "Not saving"

    # Okay, now the scheme is in place. We need to tell the system to use it:
    ifupdown(True)

# main
if __name__ == "__main__" :
    from optparse import OptionParser
    usage = """Usage: %prog [-c [-s]] [scheme] | -l | -a | -r
%prog changes your network (wi-fi) settings in /etc/network/interfaces.
You can make named schemes, like "home" or "work", matching places you
go frequently, or just make temporary schemes for places you visit briefly.
"""
    versionstr = "%prog 0.6: Set wireless network schemes.\n\
Copyright 2010 by Akkana Peck; share and enjoy under the GPL v.2 or later."
    parser = OptionParser(usage=usage, version=versionstr)
    parser.add_option("-l", "--list",
                      action="store_true", dest="list_schemes", default=False,
                      help="List the known schemes")
    parser.add_option("-a", "--accesspoints",
                      action="store_true", dest="list_accesspoints",
                      default=False,
                      help="List available accesspoints")
    parser.add_option("-r", "--reset",
                      action="store_true", dest="reset_current_scheme",
                      default=False,
                      help="Reset the connection without changing the scheme")
    parser.add_option("-s", "--save",
                      action="store_true", dest="save_new_scheme",
                      default=False,
                      help="Reset the connection without changing the scheme")
    parser.add_option("-c", "--create", metavar="new-SSID",
                      action="store", dest="newscheme",
                      help="Create a new scheme (temporary unless -s is also set)")
    (options, args) = parser.parse_args()

    if (options.list_schemes) :
        list_schemes()
    elif (options.list_accesspoints) :
        list_accesspoints(10)
    elif (options.reset_current_scheme) :
        reset_current_scheme()
    elif (options.newscheme) :
        print "Calling create_new_scheme", options.newscheme, options.save_new_scheme
        create_new_interfaces_file(options.newscheme, options.save_new_scheme)
    elif len(args) < 1 :
        #print parser.usage
        print_current_scheme()
    else :
        set_scheme(args[0])
        